package org.fujay.langchain4jserver.infrastructure.knowledge;

import cn.hutool.core.text.StrPool;
import cn.hutool.core.util.ObjectUtil;
import cn.hutool.core.util.StrUtil;
import dev.langchain4j.model.embedding.EmbeddingModel;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.fujay.commons.core.util.JsonUtils;
import org.fujay.commons.langchain4j.param.EmbedParam;
import org.fujay.commons.langchain4j.rag.processor.result.DocumentProcessResult;
import org.fujay.commons.langchain4j.server.DocumentService;
import org.fujay.langchain4jserver.domain.knowledge.aggregate.KnowledgeBaseAggregate;
import org.fujay.langchain4jserver.domain.knowledge.entity.File;
import org.fujay.langchain4jserver.domain.knowledge.exception.KnowledgeBaseNotFoundException;
import org.fujay.langchain4jserver.domain.knowledge.exception.KnowledgeFileNotFoundException;
import org.fujay.langchain4jserver.domain.knowledge.repository.FileRepository;
import org.fujay.langchain4jserver.domain.knowledge.repository.KnowledgeBaseRepository;
import org.fujay.langchain4jserver.domain.knowledge.service.IndexingService;
import org.fujay.langchain4jserver.domain.knowledge.vo.FileId;
import org.fujay.langchain4jserver.domain.knowledge.vo.IndexStatus;
import org.fujay.langchain4jserver.domain.knowledge.vo.KnowledgeBaseId;
import org.fujay.langchain4jserver.domain.model.aggregate.Model;
import org.fujay.langchain4jserver.domain.model.exception.ModelNotFoundException;
import org.fujay.langchain4jserver.domain.model.repository.ModelRepository;
import org.fujay.langchain4jserver.domain.model.service.ModelProviderService;
import org.springframework.stereotype.Service;

import java.util.List;
import java.util.stream.Collectors;

@Slf4j
@Service
@RequiredArgsConstructor
public class IndexingServiceImpl implements IndexingService {

    private final KnowledgeBaseRepository knowledgeBaseRepository;
    private final FileRepository fileRepository;
    private final ModelProviderService modelProviderService;
    private final ModelRepository modelRepository;
    private final DocumentService documentService;

    /**
     * 异步处理单个文件的索引。 包括文档加载、切分、向量化和存储。
     *
     * @param fileId 要处理的文件 ID
     * @return DocumentProcessResult
     */
    @Override
    public DocumentProcessResult indexFile(FileId fileId) {

        File file = fileRepository
                        .findById(fileId)
                        .orElseThrow(() -> new KnowledgeFileNotFoundException(fileId));

        KnowledgeBaseAggregate aggregate =
                knowledgeBaseRepository
                        .findById(file.getKnowledgeBaseId())
                        .orElseThrow(() -> new KnowledgeBaseNotFoundException(file.getKnowledgeBaseId()));

        Model model =
                modelRepository
                        .findById(aggregate.getEmbeddingModelId())
                        .orElseThrow(() -> new ModelNotFoundException(aggregate.getEmbeddingModelId()));

        EmbeddingModel embeddingModel = modelProviderService.getEmbeddingModel(model);

        EmbedParam embedParam = new EmbedParam();
        embedParam.setFilePath(file.getFilePath());
        embedParam.setStorageStrategy(file.getStorageStrategy().getCode());
        embedParam.setUserId("");
        embedParam.setEmbeddingModel(embeddingModel);
        embedParam.setIndexName(aggregate.getIndexName());
        DocumentProcessResult result = documentService.indexDoc(embedParam);

        if (result.getFailureCount() != 0) {
            aggregate.updateFileIndexStatus(fileId, List.of(), IndexStatus.FAILED, "文件索引失败");
            log.error(JsonUtils.toJsonString(result.getErrors()));
        } else {
            aggregate.updateFileIndexStatus(
                    fileId, result.getDocumentIds(), IndexStatus.COMPLETED, StrUtil.EMPTY);
        }
        knowledgeBaseRepository.save(aggregate);
        return result;
    }

    /**
     * 异步处理整个知识库的索引。 遍历知识库下所有状态为 PENDING 或 FAILED 的文件进行索引。
     *
     * @param knowledgeBaseId 要处理的知识库 ID
     */
    @Override
    public void indexKnowledgeBase(KnowledgeBaseId knowledgeBaseId) {

        KnowledgeBaseAggregate aggregate =
                knowledgeBaseRepository
                        .findById(knowledgeBaseId)
                        .orElseThrow(() -> new KnowledgeBaseNotFoundException(knowledgeBaseId));

        List<FileId> fileToIndex =
                aggregate.getFiles().values().stream()
                        .filter(
                                file ->
                                        ObjectUtil.equal(file.getIndexStatus(), IndexStatus.PENDING)
                                                || ObjectUtil.equal(file.getIndexStatus(), IndexStatus.FAILED))
                        .map(File::getId)
                        .toList();

        fileToIndex.forEach(this::indexFile);
    }

    /**
     * 异步重建知识库索引。 可能包括清空现有索引，然后重新索引所有文件。
     *
     * @param knowledgeBaseId 要重建索引的知识库 ID
     */
    @Override
    public void rebuildIndex(KnowledgeBaseId knowledgeBaseId) {

        KnowledgeBaseAggregate aggregate =
                knowledgeBaseRepository
                        .findById(knowledgeBaseId)
                        .orElseThrow(() -> new KnowledgeBaseNotFoundException(knowledgeBaseId));
        // 清空知识库索引和向量
        deleteKnowledgeBaseIndex(knowledgeBaseId);

        aggregate.getFiles().values().stream()
                .map(File::getId)
                .forEach(this::indexFile);
    }

    /**
     * 删除指定知识库的整个向量索引。
     *
     * @param knowledgeBaseId
     */
    @Override
    public void deleteKnowledgeBaseIndex(KnowledgeBaseId knowledgeBaseId) {

        KnowledgeBaseAggregate aggregate =
                knowledgeBaseRepository
                        .findById(knowledgeBaseId)
                        .orElseThrow(() -> new KnowledgeBaseNotFoundException(knowledgeBaseId));

        List<String> documentIds = aggregate.getFiles().values().stream()
                .map(File::getDocumentId)
                .map(s->StrUtil.split(StrPool.COMMA,s))
                .flatMap(List::stream)
                .collect(Collectors.toList());

        Model model = modelRepository
                        .findById(aggregate.getEmbeddingModelId())
                        .orElseThrow(() -> new ModelNotFoundException(aggregate.getEmbeddingModelId()));

        EmbeddingModel embeddingModel = modelProviderService.getEmbeddingModel(model);

        EmbedParam embedParam = new EmbedParam();
        embedParam.setEmbeddingModel(embeddingModel);
        embedParam.setIndexName(aggregate.getIndexName());
        embedParam.setDocumentIds(documentIds);
        embedParam.setIndex(true);
        documentService.deleteDoc(embedParam);

    }

    /**
     * 从向量索引中删除与特定文件关联的所有文档段（Segments/Embeddings）。
     *
     * @param knowledgeBaseId
     * @param fileId          要删除文档段的文件 ID
     */
    @Override
    public void deleteDocumentSegments(KnowledgeBaseId knowledgeBaseId, FileId fileId) {

        KnowledgeBaseAggregate aggregate =
                knowledgeBaseRepository
                        .findByIdShallow(knowledgeBaseId)
                        .orElseThrow(() -> new KnowledgeBaseNotFoundException(knowledgeBaseId));

        File file = fileRepository.findById(fileId).orElseThrow(() -> new KnowledgeFileNotFoundException(fileId));

        List<String> documentIds = StrUtil.split(StrPool.COMMA,file.getDocumentId());

        Model model = modelRepository
                .findById(aggregate.getEmbeddingModelId())
                .orElseThrow(() -> new ModelNotFoundException(aggregate.getEmbeddingModelId()));

        EmbeddingModel embeddingModel = modelProviderService.getEmbeddingModel(model);

        EmbedParam embedParam = new EmbedParam();
        embedParam.setEmbeddingModel(embeddingModel);
        embedParam.setIndexName(aggregate.getIndexName());
        embedParam.setDocumentIds(documentIds);
        embedParam.setIndex(false);
        documentService.deleteDoc(embedParam);
    }
}
